import torchvision
from torch import nn
from torch.utils.data import DataLoader
import torch
from torchvision.models import *
from unet import unet_model
import torchvision.transforms as transforms
import matplotlib.pyplot as plt

# 调整输入图像大小为 224x224
# transform = transforms.Compose([
#     transforms.Resize((224, 224)),  # 调整输入形状为 32x32
#     # transforms.ToTensor(),  # 将数据转换为张量
#     transforms.Normalize((0.5,), (0.5,))  # 将数据归一化到[-1, 1]范围内
# ])
#
# input = torch.ones((1, 3, 91, 91))  # 假设这是你的输入图像
# resized_input_image = transform(input)
# unet_m1 = unet_model.UNet(1)
# output = unet_m1(resized_input_image)
# print(output.shape)


def s2_model_train(unet, epoch, train_dataloader, loss_fn, optimizer, device):
    total_train_step = 0
    outputs_save = []
    best_loss = float('inf')
    for i in range(epoch):
        print("第{}轮训练开始".format(i + 1))
        # 训练步骤
        running_loss = 0.0
        for j, data in enumerate(train_dataloader):
            imgs, targets = data  # img.shape=(3,224,224) targets.shape=(1,224,224)
            imgs = imgs.to(device)
            targets = targets.to(device)
            outputs = unet(imgs)    # outputs.shape=(1,224,224)
            if i == epoch - 1:
                outputs_save.append(outputs)
            targets = targets.to(torch.float32)
            loss = loss_fn(outputs, targets)  # 训练的损失
            # 优化模型
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            running_loss += loss.item()

        epoch_loss = running_loss / len(train_dataloader)
        if epoch_loss < best_loss:  # 保存loss最小的模型参数
            best_loss = epoch_loss
            torch.save(unet.state_dict(), 'best_model2.pth')
        total_train_step += 1
        print("网络2训练次数:{},loss:{}".format(total_train_step, epoch_loss))


def s2_model_test(unet, loss_fn, test_dataloader, device):
    best_net = unet
    best_net.load_state_dict(torch.load("best_model2.pth"))
    with torch.no_grad():
        for j, data in enumerate(test_dataloader):
            imgs, targets = data
            imgs = imgs.to(device)
            targets = targets.to(device)
            outputs = best_net(imgs)    # (1,1,224,224)

            # loss = loss_fn(outputs, targets)

    show_output(outputs)
    #         total_test_loss = total_test_loss + loss * 1
    # total_test_loss = total_test_loss / 684
    # print("网络2整体测试集上的loss:{}".format(total_test_loss.item()))


def show_output(outputs):
    transform = transforms.Resize((91, 91))
    de_nosie_matrix = transform(outputs.detach().squeeze(0))    # (1,91,91)
    plt.subplot(2, 2, 2)
    plt.imshow(de_nosie_matrix.squeeze(0).numpy(), cmap='viridis')
    plt.title("after s2_model")
    plt.colorbar()
    # plt.show()

